import random
import threading
import time
from collections import defaultdict

import matplotlib.pyplot as plt
from tqdm import tqdm

BOARD_SIZE = 50
ENTITY_SIZE = 1
GRAVITY = 1
ELASTICITY = 0.8
FPS = 30
SLOWMO = 1
VEL_RATE = 1
SEED = 42
PREY_OF = {
    "rock": "scissors",
    "paper": "rock",
    "scissors": "paper",
}
HUNTER_OF = {v: k for k, v in PREY_OF.items()}

random.seed(SEED)


def abs(x):
    if x < 0:
        return -x
    return x


def rps_battle(etype1, etype2):
    # determine the winner of the battle
    if etype1 == etype2:
        return 0
    if etype2 == PREY_OF[etype1]:
        return 1
    return -1


def do_simple_collision_detection(positions, esize):
    # simple collision detection: all to all, square bounding boxes
    collisions = []
    for i, pos1 in enumerate(positions):
        for j, pos2 in enumerate(positions):
            if i >= j:
                continue
            if (
                pos1[0] - esize < pos2[0] + esize
                and pos1[0] + esize > pos2[0] - esize
                and pos1[1] - esize < pos2[1] + esize
                and pos1[1] + esize > pos2[1] - esize
            ):
                collisions.append((i, j))
    return collisions


class Entity:
    def __init__(self, etype, pos, vel):
        self.etype = etype
        self.pos = pos
        self.vel = vel

    def __str__(self):
        return f"{self.etype} at {self.pos}"


class Game:
    def __init__(self, n_rocks, n_papers, n_scissors):
        self.n_rocks = n_rocks
        self.n_papers = n_papers
        self.n_scissors = n_scissors
        self.board_size = BOARD_SIZE
        self.board_bbox = (0, 0, self.board_size, self.board_size)
        self.entities = {}
        self.entity_idx = 0
        self.step_idx = 0
        self.winner = None
        self.place_entities()

    def count_entities(self):
        nrocks = npapers = nscissors = 0
        for entity in self.entities.values():
            if entity.etype == "rock":
                nrocks += 1
            if entity.etype == "paper":
                npapers += 1
            if entity.etype == "scissors":
                nscissors += 1
        return nrocks, npapers, nscissors

    def summary(self):
        nrocks, npapers, nscissors = self.count_entities()
        return f"Step {self.step_idx}: Rocks={nrocks}, Papers={npapers}, Scissors={nscissors}"

    def winning(self):
        nrocks, npapers, nscissors = self.count_entities()
        if nrocks == 0 and npapers == 0:
            return "scissors"
        if nrocks == 0 and nscissors == 0:
            return "paper"
        if npapers == 0 and nscissors == 0:
            return "rock"

    def place_entities(self):
        # random place entities within the board
        for i in range(self.n_rocks):
            self.entities[self.entity_idx] = Entity(
                etype="rock",
                pos=(
                    random.random() * self.board_size,
                    random.random() * self.board_size,
                ),
                vel=(random.random() * 10 * VEL_RATE, random.random() * 10 * VEL_RATE),
            )
            self.entity_idx += 1
        for i in range(self.n_papers):
            self.entities[self.entity_idx] = Entity(
                etype="paper",
                pos=(
                    random.random() * self.board_size,
                    random.random() * self.board_size,
                ),
                vel=(random.random() * 10 * VEL_RATE, random.random() * 10 * VEL_RATE),
            )
            self.entity_idx += 1
        for i in range(self.n_scissors):
            self.entities[self.entity_idx] = Entity(
                etype="scissors",
                pos=(
                    random.random() * self.board_size,
                    random.random() * self.board_size,
                ),
                vel=(random.random() * 10 * VEL_RATE, random.random() * 10 * VEL_RATE),
            )
            self.entity_idx += 1

    def get_centroids(self):
        pos_list = defaultdict(list)
        for entity in self.entities.values():
            pos_list[entity.etype].append(entity.pos)
        centroids = {}
        for etype, pos in pos_list.items():
            x = sum([p[0] for p in pos]) / len(pos)
            y = sum([p[1] for p in pos]) / len(pos)
            centroids[etype] = (x, y)
        return centroids

    def do_hunt(self):
        # apply velocity modification to each entity:
        # towards the centroid of its preys
        # and away from the centroid of its hunters
        centroids = self.get_centroids()
        for entity in self.entities.values():
            if PREY_OF[entity.etype] in centroids:
                dx = centroids[PREY_OF[entity.etype]][0] - entity.pos[0]
                dy = centroids[PREY_OF[entity.etype]][1] - entity.pos[1]
                entity.vel = (
                    entity.vel[0] + dx * GRAVITY / FPS,
                    entity.vel[1] + dy * GRAVITY / FPS,
                )
            if HUNTER_OF[entity.etype] in centroids:
                dx = centroids[HUNTER_OF[entity.etype]][0] - entity.pos[0]
                dy = centroids[HUNTER_OF[entity.etype]][1] - entity.pos[1]
                entity.vel = (
                    entity.vel[0] - dx * GRAVITY / FPS,
                    entity.vel[1] - dy * GRAVITY / FPS,
                )

    def do_wall_bounce(self, eidx):
        # calculate the new velocity of the entity after bouncing off the wall
        entity = self.entities[eidx]
        if (entity.pos[0] - ENTITY_SIZE < self.board_bbox[0] and entity.vel[0] < 0) or (
            entity.pos[0] + ENTITY_SIZE > self.board_bbox[2] and entity.vel[0] > 0
        ):
            entity.vel = (-entity.vel[0] * ELASTICITY, entity.vel[1])
        if (entity.pos[1] - ENTITY_SIZE < self.board_bbox[1] and entity.vel[1] < 0) or (
            entity.pos[1] + ENTITY_SIZE > self.board_bbox[3] and entity.vel[1] > 0
        ):
            entity.vel = (entity.vel[0], -entity.vel[1] * ELASTICITY)

    def do_collision_bounce(self, eidx1, eidx2):
        e1, e2 = self.entities[eidx1], self.entities[eidx2]
        # if delta x and delta v are in the same direction, no bounce, otherwise bounce
        dx = e1.pos[0] - e2.pos[0]
        dvx = e1.vel[0] - e2.vel[0]
        dy = e1.pos[1] - e2.pos[1]
        dvy = e1.vel[1] - e2.vel[1]
        if dx * dvx < 0:
            # exchange x vel
            e1.vel, e2.vel = (
                (e2.vel[0] * ELASTICITY, e1.vel[1]),
                (e1.vel[0] * ELASTICITY, e2.vel[1]),
            )

        if dy * dvy < 0:
            # exchange y vel
            e1.vel, e2.vel = (
                (e1.vel[0], e2.vel[1] * ELASTICITY),
                (e2.vel[0], e1.vel[1] * ELASTICITY),
            )

    def do_battle(self, eidx1, eidx2):
        # determine the winner of the battle
        result = rps_battle(self.entities[eidx1].etype, self.entities[eidx2].etype)
        if result == 0:
            return
        if result == 1:
            self.entities[eidx2] = Entity(
                etype=self.entities[eidx1].etype,
                pos=self.entities[eidx2].pos,
                vel=self.entities[eidx2].vel,
            )
            return
        if result == -1:
            self.entities[eidx1] = Entity(
                etype=self.entities[eidx2].etype,
                pos=self.entities[eidx1].pos,
                vel=self.entities[eidx1].vel,
            )
            return

    def get_collisions(self):
        collisions = do_simple_collision_detection(
            [v.pos for k, v in self.entities.items()],
            esize=ENTITY_SIZE,
        )
        all_eidx = list(self.entities.keys())
        collision_eidx = [(all_eidx[i], all_eidx[j]) for i, j in collisions]
        return collision_eidx

    def do_all_wall_bounces(self):
        for eidx in self.entities.keys():
            self.do_wall_bounce(eidx)

    def collide_step_end_hook(self):
        # print(self.summary())
        pass

    def move_step_end_hook(self):
        pass

    def collide_step(self):
        self.do_all_wall_bounces()
        collision_eidx = self.get_collisions()
        for eidx1, eidx2 in collision_eidx:
            self.do_battle(eidx1, eidx2)
            self.do_collision_bounce(eidx1, eidx2)
        self.collide_step_end_hook()

    def move_step(self):
        for entity in self.entities.values():
            entity.pos = (
                entity.pos[0] + entity.vel[0] * 1 / FPS,
                entity.pos[1] + entity.vel[1] * 1 / FPS,
            )
        self.move_step_end_hook()

    def step(self):
        self.collide_step()
        self.do_hunt()
        self.move_step()
        self.step_idx += 1

    def start(self):
        self.thread = threading.Thread(target=self.mainloop)
        self.stop_event = threading.Event()
        self.thread.start()
        return self.thread

    def mainloop_step(self):
        self.step()
        if winner := self.winning():
            self.winner = winner
            self.stop_event.set()
        time.sleep(1 / FPS * SLOWMO)

    def mainloop(self):
        # Interruptable game loop
        while not self.stop_event.is_set():
            self.mainloop_step()

    def stop(self):
        self.stop_event.set()
        self.thread.join()


class Renderer:
    COLORS = {
        # Periwinkle, pink, and lime
        # HEX Codes: #678CEC, #D49BAE, #BBCB50
        "rock": "#678CEC",
        "paper": "#D49BAE",
        "scissors": "#BBCB50",
    }
    BG_COLOR = "white"
    TXT_COLOR = "black"

    def __init__(self, game: Game):
        self.game = game
        self.fig, self.ax = plt.subplots()
        self.ax.set_aspect("equal")
        self.ax.set_xticks([])
        self.ax.set_yticks([])
        self.ax.set_facecolor(self.BG_COLOR)
        self.ax.set_xlim(-ENTITY_SIZE, self.game.board_size + ENTITY_SIZE)
        self.ax.set_ylim(-ENTITY_SIZE, self.game.board_size + ENTITY_SIZE)
        self.rock_squares = []
        self.paper_squares = []
        self.scissors_squares = []
        self.rock_text = self.ax.text(0, 0, "Rocks", color=self.COLORS["rock"])
        self.paper_text = self.ax.text(0, 0, "Papers", color=self.COLORS["paper"])
        self.scissors_text = self.ax.text(
            0, 0, "Scissors", color=self.COLORS["scissors"]
        )
        self.winner_text = self.ax.text(0, 0, "", color=self.TXT_COLOR)
        self.step_text = self.ax.text(0, 0, "", color=self.TXT_COLOR)
        self.update()

    def update(self):
        nrocks, npapers, nscissors = self.game.count_entities()
        self.rock_text.set_text(f"Rocks: {nrocks}")
        self.paper_text.set_text(f"Papers: {npapers}")
        self.scissors_text.set_text(f"Scissors: {nscissors}")
        self.winner_text.set_text(f"Winner: {self.game.winner}")
        self.step_text.set_text(f"Steps: {self.game.step_idx}")
        self.rock_text.set_position(
            (self.ax.get_xlim()[1], self.ax.get_ylim()[1] * 0.1)
        )
        self.paper_text.set_position(
            (self.ax.get_xlim()[1], self.ax.get_ylim()[1] * 0.3)
        )
        self.scissors_text.set_position(
            (self.ax.get_xlim()[1], self.ax.get_ylim()[1] * 0.5)
        )
        self.winner_text.set_position(
            (self.ax.get_xlim()[1], self.ax.get_ylim()[1] * 0.7)
        )
        self.step_text.set_position(
            (self.ax.get_xlim()[1], self.ax.get_ylim()[1] * 0.9)
        )
        for square in self.rock_squares:
            square.remove()
        for square in self.paper_squares:
            square.remove()
        for square in self.scissors_squares:
            square.remove()
        self.rock_squares = []
        self.paper_squares = []
        self.scissors_squares = []
        for entity in self.game.entities.values():
            square = plt.Rectangle(
                (entity.pos[0] - ENTITY_SIZE, entity.pos[1] - ENTITY_SIZE),
                ENTITY_SIZE * 2,
                ENTITY_SIZE * 2,
                color=self.COLORS[entity.etype],
                alpha=0.8,
            )
            self.ax.add_artist(square)
            if entity.etype == "rock":
                self.rock_squares.append(square)
            if entity.etype == "paper":
                self.paper_squares.append(square)
            if entity.etype == "scissors":
                self.scissors_squares.append(square)
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def start(self):
        self.fig.show()
        time.sleep(10)
        while not self.game.winner:
            self.game.step()
            if winner := self.game.winning():
                self.game.winner = winner
            self.update()
            time.sleep(1 / FPS * SLOWMO)


def main():
    game = Game(n_rocks=10, n_papers=10, n_scissors=10)
    game.start()


def test_balance():
    # 1000 games with 10 rocks, 10 papers, 10 scissors results:
    # Rock wins: 328
    # Paper wins: 333
    # Scissors wins: 339
    # Average steps: 3859.379
    global SLOWMO
    SLOWMO = 0
    winners = []
    steps = []
    for _ in tqdm(range(1000)):
        game = Game(n_rocks=10, n_papers=10, n_scissors=10)
        thread = game.start()
        thread.join()
        assert game.winner is not None
        winners.append(game.winner)
        steps.append(game.step_idx)
    print(f"Rock wins: {winners.count('rock')}")
    print(f"Paper wins: {winners.count('paper')}")
    print(f"Scissors wins: {winners.count('scissors')}")
    print(f"Average steps: {sum(steps) / len(steps)}")


def test_renderer():
    game = Game(n_rocks=20, n_papers=20, n_scissors=20)
    renderer = Renderer(game)
    renderer.start()
    while True:
        time.sleep(60)


if __name__ == "__main__":
    # main()
    # test_balance()
    test_renderer()
